#include "force.h"
#include "para.h"
#include "physical_properties.h"
#include<vector>
#include<iostream>

/**************************************************************************Force*****************************************************************************/
void Force::distance(const std::vector<double>& ri, const std::vector<double>& rj, double& dx, double& dy, double& dz, double& r) {
    dx = ri[0] - rj[0];
    dy = ri[1] - rj[1];
    dz = ri[2] - rj[2];

    if (dx > Parameter_Box_size / 2.0)
        dx = dx - Parameter_Box_size;
    if (dx < -Parameter_Box_size / 2.0)
        dx = dx + Parameter_Box_size;
    if (dy > Parameter_Box_size / 2.0)
        dy = dy - Parameter_Box_size;
    if (dy < -Parameter_Box_size / 2.0)
        dy = dy + Parameter_Box_size;
    if (dz > Parameter_Box_size / 2.0)
        dz = dz - Parameter_Box_size;
    if (dz < -Parameter_Box_size / 2.0)
        dz = dz + Parameter_Box_size;
    double r2 = dx * dx + dy * dy + dz * dz;
    r = sqrt(r2);
    return;
}


/**************************************************************************LJ_Force**************************************************************************/
void LJ_Force::terms(const double& r, double& attractive_term, double& repulsive_term) {
    attractive_term = pow((Constant_Sigma / r), 6);
    repulsive_term = pow(attractive_term, 2);
    return;
}
double LJ_Force::force(const double& r, const double& attractive_term, const double& repulsive_term) {
    double fij = 4.0 * Constant_Epsilon * (12.0 * repulsive_term - 6.0 * attractive_term) / pow(r, 2);
    return fij;
}
double LJ_Force::energy(const double& attractive_term, const double& repulsive_term) {
    double potential_energy = 4.0 * Constant_Epsilon * (repulsive_term - attractive_term);
    return potential_energy;
}
void LJ_Force::Compute_Forces(std::vector<Atom>& mol, double& potential_energy, double& virial){
    potential_energy = 0.0;
    virial = 0.0;
    int Natoms = mol.size();
    double dx = 0, dy = 0, dz = 0, r = 0, r2 = 0, attractive_term = 0, repulsive_term = 0;

    // initialize force
    for (int i = 0; i < Natoms; i++)
        for (int j = 0; j < 3; j++)
            mol[i].force[j] = 0;

    for (int i = 0; i < Natoms - 1; i++) {
        for (int j = i + 1; j < Natoms; j++) {
            LJ_Force::distance(mol[i].coord, mol[j].coord, dx, dy, dz, r);
            if (r2 < Parameter_Rcutoff) {
                LJ_Force::terms(r, attractive_term, repulsive_term);
                double fij = LJ_Force::force(r, attractive_term, repulsive_term);
                mol[i].force[0] += fij * dx;
                mol[j].force[0] -= fij * dx;
                mol[i].force[1] += fij * dy;
                mol[j].force[1] -= fij * dy;
                mol[i].force[2] += fij * dz;
                mol[j].force[2] -= fij * dz;

                double energyij = LJ_Force::energy(attractive_term, repulsive_term); 
                potential_energy = potential_energy + energyij;
                virial = virial + 4.0 * Constant_Epsilon * (12.0 * repulsive_term - 6.0 * attractive_term);
            }
        }
    }

    virial = virial / (3.0 * pow(Parameter_Box_size, 3)); 

    return;
}

/***********************************PIMD***********************************/
void LJ_Force::Compute_Forces_PIMD(std::vector< std::vector<Atom> >& mirr, std::vector<double>& rdf, std::vector<double>& potential_energy, std::vector<double>& virial) {
    for (int k = 0; k < Parameter_nbeads; k++)
        LJ_Force::Compute_Forces(mirr[k], potential_energy[k], virial[k]);
}

std::vector< std::vector<double> > LJ_Force::Compute_Mirror_Center(std::vector< std::vector<Atom> >& mirr) {
    vector< vector<double> > Mirror_Center(Parameter_Natoms, vector<double>(3, 0));

    for (int j = 0; j < 3; j++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int k = 0; k < Parameter_nbeads; k++) {
                Mirror_Center[i][j] += mirr[k][i].coord[j];
            }
            Mirror_Center[i][j] /= Parameter_nbeads;
        }
    }
    return Mirror_Center;
}

